Vertex AIのバッチ予測で入出力の対象をBigQueryにする

Google Cloud AI/MLブログリレー2023

#Google Cloud (GCP)

#Vertex AI

#Google BigQuery

#機械学習

じょんすみす

2023.12.10

Vertex AIのバッチ予測で入出力の対象をBigQueryにする

おはこんハロチャオ～！何者(なにもん)なんじゃ？じょんすみすです。

Vertex AIではModel Registryに登録されたモデルを使ってバッチ推論を実行できます。バッチ推論を行うために、必要な設定やCloud Storageを利用する際の入力に関しては以前紹介しました。

Vertex AIのカスタムコンテナでバッチ推論をする | DevelopersIO

今回は入力や出力をBigQueryとして動かしてみます。

Model Registryとエンドポイント, バッチ推論

まずは、簡単に機能をおさらいしておきましょう。

Vertex AIではModel Registryに以下の2つの情報を登録しておきます。

モデルのCloud Storage上のパス
推論処理を実行するためのコンテナのパス
コンテナ実行時のコマンド

これらの情報に基づいて、エンドポイントを立てる際にはコンテナを起動し、その中から処理を実施する仕組みになっています。この際、コンテナ内には推論をするために処理が実装されたコードを含めますが、HTTPでリクエストを受け付ける仕組みになっている必要があります。環境変数 AIP_HEALTH_ROUTE, AIP_PREDICT_ROUTE で渡されるパスへのそれぞれGETとPOSTでリクエストを受け付ける仕組みを実装しておけばいいわけです。推論に利用するモデルのパスも環境変数 AIP_STORAGE_URI として渡されるため、 AIP_PREDICT_ROUTE へのリクエストがあった際に推論結果を返すコードを実装します。

リクエストはJSON形式で、 instances というキーに対して配列で複数の値が渡されます。これらが推論対象となるデータです。推論結果も同様にJSON形式で返します。 predictions というキーの中なに入力で与えられたそれぞれに対する推論結果を配列で返します。

これらの配列の中身自体は、どのような形式で与えるかは仕様としては決まっていませんので、JSONで記述可能な任意の形式で与えることが可能です。よくある形式としては、入力は特徴を列挙した配列や特徴量名と値のkey-value形式となるでしょう。出力に関しても、推論結果だけを返す場合や、それぞれの予測ラベルに対する確信度を返すなどが考えられます。

バッチ推論でも同様にModel Registryに登録されているモデルを使って推論を実施します。バッチ推論の入出力はCloud StorageやBigQueryが入出力となります。これらは一致している必要はなく、入力はBigQueryで出力はCloud Storageとしたりその逆にすることも可能です。

バッチ推論でも同じモデルを利用することから、内部ではHTTPによるリクエストで処理を行いますが、利用者が意識しなくて以下の部分を実施してくれます。

入力元からデータを取得
エンドポイントにリクエスト投げて推論結果を取得
取得した結果を出力先に書き込む

これらをVertex AI側でやってくれているため、同じモデルでエンドポイントの作成とバッチ推論の両方を行えるわけです。

ただし、自動的にやってくれているため、利用者側で前後の処理を挟むことはできません。そのため、特に入力に関してはバッチ推論の仕様に合わせて内部の処理が実行できるように実装しておく必要があります。

BigQueryを入出力として利用する

BigQueryを入出力として利用する方法はそれほど難しくありません。

コンソール上で行う時は、以下のようにそれぞれでBigQueryを選択して対象のテーブルを指定します。

また、Pyhton SDKから呼び出す際も引数でBigQuery用のものを指定するのみとなります。

# aiplatform.Model("<モデルID>")で既存のモデルを取得するか
# 学習ジョブのrunの結果を受け取る
model = ...

# 入力はbigquery_source、出力はbigquery_destination_prefixで指定
# gcs_source, gcs_distination_prefixは指定しない
model.batch_predict(
    job_display_name="titanic-predict",
    bigquery_source="bq://<project id>.<dataset>.titanic_test",
    bigquery_destination_prefix="bq://<project id>.<dataset>.titanic_batch_predict",
    machine_type="n1-standard-4"
)

これだけで利用が可能ですが、注意点としては入出力の仕様となります。

入力は1レコードごとにカラム名をkey、その値をvalueとした形式で、instancesの中身がその配列になっているような形式となります。一方、出力はpredictionの中身が1レコードごとに順に付与されます。

以下のように返ってきたJSONが付与されたテーブルが出力されます。 (Titanicデータセットに対して画面表示用にいくつかの特徴量に絞ったもので推論を実施しています)

上述の記載の通り、BigQueryからの入力はテーブル名を指定する形式になっています。そのため、SELECT文を使って何らかの処理を行ったものを渡すということができません。公式に提供されているサンプルでは、リクエスト時に対象とするカラムや除外するカラムを設定する例が記載されていますが、なんらかの特徴量エンジニアリングを行うなどはここではできません。

これを踏まえ、他での実用性の観点からも見ると推論時の処理として学習時と同じ特徴量エンジニアリングを行う部分を内部のコードで実装しておくといいでしょう。